package com.likeycy.my.tree;

import java.util.Comparator;

/**
 * @ClassName: AVLTree
 * @Description: 实现Java版本的AVL树
 * @Author: sodagreen
 * @Date: 2021/3/18 23:58
 * @Version: 1.0
 */
public class AVLTree<E> extends MyBinarySearchTree<E> {

    public AVLTree() {
        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    /*
     * 问题一、添加导致的失衡
     * 最坏情况：可能会导致所有祖先节点失衡
     * 父节点、非祖父节点，都不会失衡
     *
     *
     */

    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡
                rebalance(node);
                // 整棵树恢复平衡
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡
                rebalance(node);
            }
        }
    }

    @Override
    protected Node<E> newNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    /**
     * 第二种实现：恢复二叉树平衡
     *
     * @param grandparentNode 高度最底的那个不平衡节点(祖父节点)
     */
    private void rebalance2(Node<E> grandparentNode) {
        AVLNode<E> avlNode = (AVLNode<E>) grandparentNode;
        Node<E> parent = avlNode.tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        // 进入if，表示失衡的左子树节点，进入else表示失衡的右子树节点
        if (parent.isLeftChild()) {
            // 进入if则是 LL ，进入else则是 LR
            if (node.isLeftChild()) {
                rotateRight(grandparentNode);
            } else {
                rotateLeft(parent);
                rotateRight(grandparentNode);
            }
        } else {
            // 进入if则是 RL ，进入else则是 LL
            if (node.isLeftChild()) {
                rotateRight(parent);
                rotateLeft(grandparentNode);
            } else {
                rotateLeft(grandparentNode);
            }
        }
    }

    /**
     * 统一旋转
     * @param grand
     */
    private void rebalance3(Node<E> grand) {
        AVLNode<E> avlNode = (AVLNode<E>) grand;
        Node<E> parent = avlNode.tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                // LL 右旋转
                avlTreeRotate(grand, node, node.right, parent, parent.right, grand);
            } else {
                // LR 双旋转：LL右旋转后再RR左旋
                avlTreeRotate(grand, parent, node.left, node, node.right, grand);
            }
        } else {
            if (node.isLeftChild()) {
                // RL 双旋转：RR左旋转后再LL右旋转
                avlTreeRotate(grand, grand, node.left, node, node.right, parent);
            } else {
                // RR 左旋转
                avlTreeRotate(grand, grand, parent.left, parent, node.left, node);
            }
        }

    }

    /**
     * 统一旋转
     * @param grand
     */
    private void rebalance(Node<E> grand) {
        AVLNode<E> avlNode = (AVLNode<E>) grand;
        Node<E> parent = avlNode.tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();
        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                // LL 右旋转
                rotate(grand, node.left, node, node.right, parent, parent.right, grand, grand.right);
            } else {
                // LR 双旋转：LL右旋转后再RR左旋
                rotate(grand, parent.left, parent, node.left, node, node.right, grand, grand.right);
            }
        } else {
            if (node.isLeftChild()) {
                // RL 双旋转：RR左旋转后再LL右旋转
                rotate(grand, grand.left, grand, node.left, node, node.right, parent, parent.right);
            } else {
                // RR 左旋转
                rotate(grand, grand.left, grand, parent.left, parent, node.left, node, node.right);
            }
        }

    }

    private void avlTreeRotate(Node<E> r,
                        Node<E> b, Node<E> c,
                        Node<E> d,
                        Node<E> e, Node<E> f) {
        d.parent = r.parent;
        if (r.isLeftChild()) {
            r.parent.left = d;
        } else if (r.isRightChild()) {
            r.parent.right = d;
        } else {
            root = d;
        }

        // b-c
        b.right = c;
        if (c != null) {
            c.parent = b;
        }
        updateHeight(b);

        // e-f
        f.left = e;
        if (e != null) {
            e.parent = f;
        }
        updateHeight(f);

        // b-d-f
        d.left = b;
        d.right = f;
        b.parent = d;
        f.parent = d;
        updateHeight(d);
    }

    /**
     * 统一旋转二叉树节点，让二叉树平衡
     * @param r 子树的根节点，旋转后 root = d;
     * @param a
     * @param b ac的根节点
     * @param c
     * @param d 该节点旋转后称为左右子树的根节点，也就是root
     * @param e
     * @param f ed的根节点
     * @param g
     */
    private void rotate(Node<E> r,
                        Node<E> a, Node<E> b, Node<E> c,
                        Node<E> d,
                        Node<E> e, Node<E> f, Node<E> g) {
        d.parent = r.parent;
        if (r.isLeftChild()) {
            r.parent.left = d;
        } else if (r.isRightChild()) {
            r.parent.right = d;
        } else {
            root = d;
        }

        // a-b-c
        b.left = a;
        if (a != null) {
            a.parent = b;
        }
        b.right = c;
        if (c != null) {
            c.parent = b;
        }
        updateHeight(b);

        // b-d-f
        f.left = e;
        if (e != null) {
            e.parent = f;
        }
        f.right = g;
        if (g != null) {
            g.parent = f;
        }
        updateHeight(f);

        // d-b-f
        d.left = b;
        d.right = f;
        b.parent = d;
        f.parent = d;
        updateHeight(d);
    }


    /**
     * 失衡节点左旋转
     * @param grandparentNode
     */
    private void rotateLeft(Node<E> grandparentNode) {
        Node<E> parent = grandparentNode.right;
        Node<E> child = parent.left;
        grandparentNode.right = child;
        parent.left = grandparentNode;

        afterRotate(grandparentNode, parent, child);
    }

    /**
     * 失衡节点右旋转
     * @param grandparentNode
     */
    private void rotateRight(Node<E> grandparentNode) {
        Node<E> parent = grandparentNode.left;
        Node<E> child = parent.right;
        grandparentNode.left = child;
        parent.right = grandparentNode;

        afterRotate(grandparentNode, parent, child);
    }

    /**
     * 旋转后维护祖父节点、父节点、子节点的parent属性，并重新更新高度
     * @param grandparentNode
     * @param parent
     * @param child
     */
    private void afterRotate(Node<E> grandparentNode, Node<E> parent, Node<E> child) {
        // 让parent成为子树的根节点
        parent.parent = grandparentNode.parent;
        if (grandparentNode.isLeftChild()) {
            grandparentNode.parent.left = parent;
        } else if (grandparentNode.isRightChild()) {
            grandparentNode.parent.right = parent;
        } else {
            // grandparentNode是root节点
            root = parent;
        }

        // 更新 child 的parent属性
        if (child != null) {
            child.parent = grandparentNode;
        }

        // 更新 grandparentNode 的 parent
        grandparentNode.parent = parent;

        // 更新高度
        updateHeight(grandparentNode);
        updateHeight(parent);
    }

    /**
     * 更新节点的高度
     *
     * @param node 新增的节点
     */
    private void updateHeight(Node<E> node) {
        AVLNode<E> avlNode = (AVLNode<E>) node;
        avlNode.updateHeight();
    }

    /**
     * 节点是否平衡
     *
     * @param node 新增的节点
     * @return
     */
    private boolean isBalanced(Node<E> node) {
        AVLNode<E> avlNode = (AVLNode<E>) node;
        // 求绝对值
        return Math.abs(avlNode.balanceFactor()) <= 1;
    }

    /**
     * 实现自己的AVL树Node节点，增加高度属性，继承BST树的Node其他属性
     * @param <E>
     */
    private static class AVLNode<E> extends Node<E> {
        /**
         * 节点高度
         */
        int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 计算左右子树的平衡因子
         * @return
         */
        public int balanceFactor() {
            // left实例为Node<E>，而非AVLNode<E>。 所以没有height属性，需要强制转型
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        /**
         * 更新节点的高度
         */
        public void updateHeight() {
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        public Node<E> tallerChild() {
            int leftHeight = (left == null) ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = (right == null) ? 0 : ((AVLNode<E>) right).height;
            if (leftHeight > rightHeight) {
                return left;
            }
            if (leftHeight < rightHeight) {
                return right;
            }

            // 高度一样则返回同方向的，左子节点则返回左，否则返回右
            return isLeftChild() ? left : right;
        }

    }
}
